Skip to content

Only call torch.autocast if it will have an effect#42747

Merged
ArthurZucker merged 5 commits intohuggingface:mainfrom
hmellor:maybe-autocast
Dec 10, 2025
Merged

Only call torch.autocast if it will have an effect#42747
ArthurZucker merged 5 commits intohuggingface:mainfrom
hmellor:maybe-autocast

Conversation

@hmellor
Copy link
Member

@hmellor hmellor commented Dec 9, 2025

Adds a context manager to replace direct use of torch.autocast that only autocasts if:

  • autocast is already enabled in this context
  • Or this call to maybe_autocast has enabled=True

This prevents autocast being added to the graph when it is effectively a no-op.
Which makes graph splitting in torch.compile more flexible as it removes the requirement that partition IDs be monotonically increasing.

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perfect thanks for pushing this!

@github-actions
Copy link
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: afmoe, apertus, arcee, aria, bamba, bitnet, blt, chameleon, cohere, cohere2, csm, cwm, dbrx, decision_transformer, deepseek_v2, deepseek_v3

@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@ArthurZucker ArthurZucker merged commit 6d0adb5 into huggingface:main Dec 10, 2025
25 checks passed
@hmellor hmellor deleted the maybe-autocast branch December 10, 2025 10:43
SangbumChoi pushed a commit to SangbumChoi/transformers that referenced this pull request Jan 23, 2026
* Only call `torch.autocast` if it will have an effect

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

* whitespace

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

* fixup

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

* fix copies

---------

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Arthur <arthur.zucker@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants